"""
The LossCCE class implements categorical cross entropy loss.
"""

import numpy as np

class LossCCE:

    def __init__(self):
        """
        No initialization necessary.
        """
        pass

    def forward_pass(self, y_pred, y_true):
        """
        Compute the categorical cross entropy loss.
        y_pred and y_true should be numpy arrays of shape (classes, batch_size),
        with y_true in one-hot format.
        """
        self.y_pred = y_pred
        self.y_true = y_true
        # Clip predictions to prevent log(0)
        self.clipped_preds = np.clip(y_pred, 1e-15, 1 - 1e-15)
        # Compute negative log-likelihood per sample
        self.sample_losses = -np.sum(y_true * np.log(self.clipped_preds), axis=0)
        return np.mean(self.sample_losses)

    def backward_pass(self):
        """
        Compute gradient of loss with respect to y_pred.
        """
        # Number of samples
        batch_size = self.y_true.shape[1]
        # Gradient on predictions
        self.dinputs = - self.y_true / self.clipped_preds
        # Normalize by batch size
        self.dinputs = self.dinputs / batch_size
        return self.dinputs